from typing import Union, Any, overload
import torch
from feasytools import ArgChecker
from xml.etree import ElementTree as ET

class WorkerArguments:
    @overload
    def __init__(self,pars:ET.Element):...
    @overload
    def __init__(self,pars:ArgChecker): ...
    @overload
    def __init__(self,pars:'dict[str,Any]'): ...

    def __init__(self,pars:'Union[ET.Element,ArgChecker,dict[str,Any]]'):
        if isinstance(pars,ET.Element): pars=ArgChecker(pars.attrib)
        elif isinstance(pars,dict): pars=ArgChecker(pars)
        # running mode
        self.mode = pars.pop_str("mode","train") # [train, eval, draw]
        if not self.mode in ['train','eval','draw']:
            raise ValueError(f"Invalid mode: {self.mode}. Can only be 'train', 'eval' or 'draw'.")
        
        # Path parameters
        self.data = pars.pop_str("data") # [traffic12, traffic12_v2g, traffic37, traffic37_v2g]
        self.path = pars.pop_str("path",f"./data/{self.data}/") # root path of the data file

        # Architecture selection.
        self.model = pars.pop_str("model") # [Pyraformer, LSTM, DNN, Transformer]
        self.decoder = pars.pop_str('decoder', 'FC') # [FC, attention]

        # Dataloader parameters.
        self.scaler = pars.pop_str("scaler","0max") # [0max, std]
        self.ilen = pars.pop_int("ilen") # 1440
        self.olen = pars.pop_int("olen") # 60
        self.xpm = pars.pop_str('xpm','sum') # [sum, mean]
        self.ypm = pars.pop_str('ypm','mean') # [sum, mean]
        self.sr = pars.pop_int("sr",1)
        self.inverse = pars.pop_bool("inverse") # denormalize output data
        self.non_overlap = pars.pop_bool("non_overlap")
        self.non_future = pars.pop_bool("non_future")

        # Training parameters.
        self.epoch = pars.pop_int('epoch', 5)
        self.batch_size = pars.pop_int('batch_size', 32)
        self.pretrain = pars.pop_bool('pretrain')
        self.hard_sample_mining = pars.pop_bool('hard_sample_mining')
        self.dropout = pars.pop_float('dropout', 0.05)
        self.lr = pars.pop_float('lr', 1e-4)
        self.lr_step = pars.pop_float('lr_step', 0.1)

        # Common Model parameters.
        self.d_model = pars.pop_int('d_model', 512)
        self.d_inner_hid = pars.pop_int('d_inner_hid', 512)
        self.d_k = pars.pop_int("d_k", 128)
        self.d_v = pars.pop_int("d_v", 128)
        self.d_bottleneck = pars.pop_int('d_bottleneck', 128)
        self.n_head = pars.pop_int("n_head", 8)
        self.n_layer = pars.pop_int("n_layer", 4)
        self.t_size = pars.pop_int("t_size", 3)

        # Pyraformer parameters.
        ws = pars.pop_str("window_size","[3,3,3]")        
        self.window_size = list(map(int,ws.strip("[]").split(","))) # The number of children of a parent node.
        self.inner_size = pars.pop_int('inner_size',3) # The number of ajacent nodes.
        # CSCM structure. selection: [Bottleneck_Construct, Conv_Construct, MaxPooling_Construct, AvgPooling_Construct]
        self.CSCM = pars.pop_str('CSCM', 'Bottleneck_Construct')
        self.truncate = pars.pop_bool('truncate') # Whether to remove coarse-scale nodes from the attention structure
        self.use_tvm = pars.pop_bool('use_tvm') # Whether to use TVM.

        # Experiment repeat times.
        self.iter_num = pars.pop_int('iter_num', 1) # Repeat number.

        if not pars.empty(): raise ValueError(f"Invalid paramerters exist: {pars.keys()}")

        assert self.ilen % self.sr == 0 and self.olen % self.sr == 0
        self.input_size = self.ilen // self.sr
        self.predict_step = self.olen // self.sr

        if self.mode == 'draw':
            self.iter_num = 1
            self.inverse = True
            self.non_overlap = True
        
        if torch.cuda.is_available():
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")
        
        # 对于无V2G的数据集，输入文件为行程信息和充电价格，输出文件为充电负荷
        # 对于有V2G的数据集，输入文件为行程信息、充电价格和放电价格，输出文件为充电负荷、V2G放电负荷和最大可能的V2G放电功率
        out_fea = {
            'traffic12':45, # 40PK + 5CS
            'traffic12_v2g':125, # 40PK*3 + 5CS
            'traffic37':140, # 130PK + 10CS
            'traffic37_v2g':400, # 130PK*3 + 10CS
        }
        in_fea = {
            'traffic12':125, # 40PK_TRIP*2 + 40PK_PRICE + 5CS_PRICE
            'traffic12_v2g':165, # 40PK_TRIP*2 + 40PK_PRICE*2 + 5CS_PRICE
            'traffic37':400, # 130PK_TRIP*2 + 130PK_PRICE + 10CS_PRICE
            'traffic37_v2g':530, # 130PK_TRIP*2 + 130PK_PRICE*2 + 10CS_PRICE
        }
        self.enc_in = in_fea[self.data] + out_fea[self.data]
        self.dec_in = out_fea[self.data]
        self.covariate_size = 3
        self.seq_num = 1
        self.embed_type = 'DataEmbedding'
    
    def __str__(self):
        return str(self.__dict__)